package org.pz.filter;

import lombok.extern.slf4j.Slf4j;
import org.pz.pojo.TokenBucketRateLimiter;
import org.springframework.cloud.gateway.filter.GatewayFilter;
import org.springframework.cloud.gateway.filter.factory.AbstractGatewayFilterFactory;
import org.springframework.http.HttpStatus;
import org.springframework.stereotype.Component;

/**
 * 限流过滤器，自定义限流
 * GatewayFilter，网关过滤器，需要通过 spring.cloud.routes.filters 配置在具体路由下，只作用在当前路由上或通过 spring.cloud.default-filters 配置在全局，作用在所有路由上
 * GatewayFilter 过滤器需要在 routes 中指定
 * */
@Component
@Slf4j
public class LimitFilter extends AbstractGatewayFilterFactory {

    // 初始化限流器（示例：桶容量=5，每 10 秒填充3个令牌）
    private final TokenBucketRateLimiter rateLimiter = new TokenBucketRateLimiter(5, 3, 10_000);

    @Override
    public GatewayFilter apply(Object config) {

        log.info("自定义的限流过滤器-------------开始");

        return (exchange, chain) -> rateLimiter.tryAcquire()
                .flatMap(allow -> {
                    if (allow) {
                        log.info("自定义的限流过滤器-------放行------开始");
                        return chain.filter(exchange);
                    } else {
                        exchange.getResponse().setStatusCode(HttpStatus.TOO_MANY_REQUESTS);
                        log.info("自定义的限流过滤器-------限流了------结束");
                        return exchange.getResponse().setComplete();
                    }
                });
    }
}
